package leecode

import (
	"math/rand"
)

func smallestK(arr []int, k int) []int {
	quickSort(arr)
	return arr[:k]
}

func quickSort(arr []int) {
	if arr == nil || len(arr) < 2 {
		return
	}
	quickSortHelper(arr, 0, len(arr)-1)
}

func quickSortHelper(arr []int, l, r int) {
	if l < r {
		pivotIndex := l + rand.Intn(r-l+1) //得到一个随机数
		arr[pivotIndex], arr[r] = arr[r], arr[pivotIndex]
		p := partition(arr, l, r)
		quickSortHelper(arr, l, p[0]-1)
		quickSortHelper(arr, p[1]+1, r)
	}
}

func partition(arr []int, l, r int) []int {
	less := l - 1 // <区的右边界
	more := r     // > 区 的左边界
	pivot := arr[r]

	for l < more {
		if arr[l] < pivot {
			less++
			arr[l], arr[less] = arr[less], arr[l]
			l++
		} else if arr[l] > pivot {
			more--
			arr[l], arr[more] = arr[more], arr[l]
		} else {
			l++
		}
	}

	arr[more], arr[r] = arr[r], arr[more]
	return []int{less + 1, more}
}
